package com.huawei.java.solver;

import java.util.ArrayList;
import java.util.List;

public class Dinic {
    private static int MAX = 1000;
    private int vertex;
    private int node; // 边缘节点的数目
    private int k; // 边的数目
    private int customer; // 客户的数量
    private Integer[][] allocation; // 节点i可以向j提供的带宽
    private List<Integer> supply; // i节点可以提供的总带宽
    private List<Integer> demand; // j客户需要的总带宽
    private boolean[] used = new boolean[MAX];
    private int[] pre = new int[MAX];
    private int[] distance = new int[MAX];
    private int[] capacity = new int[MAX];
    private List<Edge>[] map; // 邻接表

    public Dinic(Integer[][] allocation, List<Integer> supply, List<Integer> demand) {
        this.node = supply.size();
        this.customer = demand.size();
        this.vertex = node + customer + 2; // 总共的图的顶点数目
        this.k = node + customer + node * customer; // 总共边的数目
        this.allocation = allocation; // 资源分配
        this.demand = demand;
        this.supply = supply;
    }

    public void init() {
        map = new ArrayList [vertex]; // 总共的节点数目为n + 2
        for (int i = 0; i < vertex; i++)
            map[i] = new ArrayList<Edge>();

        // 规定源节点S的编号为0
        // 规定边缘节点的编号为1 --> n, 客户节点的编号为n + 1 --> n + customer
        // 规定汇节点T的编号为 n + customer + 1

        // 1. 添加S->边缘节点的边
        for (int i = 1; i <= node; i++) {
            int from = 0, to = i;
            int cap = supply.get(i - 1), cost = 1;
            map[from].add(new Edge(from, to, cap, cost));
            map[to].add(new Edge(to, from, 0, -cost));
        }

        // 2. 添加客户节点->T的边
        for (int i = node + 1; i <= node + customer; i++) {
            int from = i, to = customer + node + 1;
            int cap = demand.get(i - node - 1), cost = 1;
            map[from].add(new Edge(from, to, cap, cost));
            map[to].add(new Edge(to, from, 0, -cost));
        }

        // 3. 添加边缘节点->客户节点的边
        for (int i = 1; i <= node; i++) {
            for (int j = node + 1; j <= node + customer; j++) {
                int from = i, to = j;
                int cap = allocation[i - 1][j - node - 1];
                int cost = 0;
                map[from].add(new Edge(from, to, cap, cost));
                map[to].add(new Edge(to, from, 0, -cost));
            }
        }
    }

    // 寻找从start到end的最短路径
    public boolean spfa(int start, int end) {
        int[] count = new int[vertex];
        for (int i = 0; i < vertex; i++) {
            used[i] = false;
            pre[i] = -1;
            distance[i] = Integer.MAX_VALUE;
            capacity[i] = Integer.MAX_VALUE;
        }
        used[start] = true;
        pre[start] = start;
        distance[start] = 0;
        count[start]++;
        List<Integer> queue = new ArrayList<>();
        queue.add(start);
        while (!queue.isEmpty()) {
            int index = queue.get(0);
            queue.remove(0);
            used[index] = false;
            for (int i = 0; i < map[index].size(); i++) {
                Edge temp = map[index].get(i);
                if (temp.getCap() > 0 && distance[temp.getTo()] > distance[index] + temp.getCost()) {
                    distance[temp.getTo()] = distance[index] + temp.getCost();
                    pre[temp.getTo()] = index;
                    capacity[temp.getTo()] = Math.min(capacity[index], temp.getCap());
                    if (!used[temp.getTo()]) {
                        used[temp.getTo()] = true;
                        queue.add(temp.getTo());
                        count[temp.getTo()]++;
                        if (count[temp.getTo()] > vertex)
                            return false;
                    }
                }
            }
        }
        if (distance[end] != Integer.MAX_VALUE && pre[end] != -1)
            return true;
        return false;
    }

    public int getResult() {
        init();
        int minCost = 0;
        int start = 0;
        int end = vertex - 1;
        while (true) {
            if (!spfa(start, end))
                break;
            minCost += distance[end] * capacity[end];
            int last = end;
            int begin = end;
            while (begin != start) {
                last = begin;
                begin = pre[last];
                int i = 0, j = 0;
                for (; i< map[begin].size(); i++) {
                    if (map[begin].get(i).getTo() == last)
                        break;
                }
                map[begin].get(i).setCap(map[begin].get(i).getCap() - capacity[end]);
                for (; j < map[last].size(); j++) {
                    if (map[last].get(j).getTo() == begin)
                        break;
                }
                map[last].get(j).setCap(map[last].get(j).getCap() + capacity[end]);
            }
        }
        return minCost;
    }

    // M[i][j]表示边缘节点i分配给客户j的带宽
    public List<List<Integer>> getPlan() {
        List<List<Integer>> res = new ArrayList<>(); // 生成结果矩阵
        for (int i = 0; i < node; i++) res.add(new ArrayList<>());
        for (int i = 0; i < node; i++) {
            for (int j = 0; j < customer; j++) {
                res.get(i).add(map[node + j + 1].get(i + 1).getCap());
            }
        }
        return res;
    }

    //M[i][j]表示客户节点i需要边缘节点j的分配
    public Integer[][] getAnotherPlan(){
        Integer[][] res = new Integer[customer][node];// 生成结果矩阵
        for (int i = 0; i < customer; ++i){
            for (int j = 0; j < node; ++j){
                res[i][j] = map[node + i + 1].get(j + 1).getCap();
            }
        }
        return res;
    }
}
